/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * Global Trust Authority is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::future::Future;
use std::pin::Pin;

use common_log::{error, info};
use policy::api::{get_policy_by_ids, query_policy};
use policy_engine::evaluate_policy;
use rdb::get_connection;

use crate::{
    chain::attestation_chain::{AttestationHandler, AttestationContext},
    entities::token::PolicyInfo,
    error::attestation_error::AttestationError,
};

/// Handler for custom policy evaluation in the attestation chain
pub struct CustomPolicyHandler {
    next: Option<Box<dyn AttestationHandler>>,
}

impl CustomPolicyHandler {
    pub fn new() -> Self {
        Self { next: None }
    }

    /// Get default policies for the given attester type
    async fn get_default_policy_ids(
        &self,
        attester_type: &str,
        user_id: &str,
    ) -> Result<Vec<String>, AttestationError> {
        info!("No policy_ids provided, using default policies for attester_type: {}", attester_type);
        
        let db_connection = get_connection().await.map_err(|e| {
            error!("Failed to get database connection: {}", e);
            AttestationError::DatabaseError
        })?;
        
        match query_policy::get_default_policies_by_type(&db_connection, attester_type.to_string(), user_id).await {
            Ok(default_policies) => {
                if !default_policies.is_empty() {
                    Ok(default_policies.iter().map(|p| p.id.clone()).collect())
                } else {
                    info!("No default policies found for attester_type: {}", attester_type);
                    Ok(Vec::new())
                }
            },
            Err(e) => {
                error!("Failed to get default policies for attester_type {}: {}", attester_type, e);
                Err(AttestationError::DatabaseError)
            },
        }
    }

    /// Build policy ID list from provided IDs or default policies
    async fn build_policy_id_list(
        &self,
        policy_ids: Option<&Vec<String>>,
        attester_type: &str,
        user_id: &str,
    ) -> Result<Vec<String>, AttestationError> {
        match policy_ids {
            Some(ids) => Ok(ids.clone()),
            None => self.get_default_policy_ids(attester_type, user_id).await,
        }
    }

    /// Validate policy constraints for specific token formats
    fn validate_policy_constraints(
        &self,
        token_fmt: &str,
        policy_count: usize,
    ) -> Result<(), AttestationError> {
        if token_fmt == "ear" && policy_count > 1 {
            return Err(AttestationError::InvalidParameter(
                "ear token only support one policy".to_string()
            ));
        }
        Ok(())
    }

    /// Evaluate user policies
    async fn evaluate_user_policies(
        &self,
        verify_evidence: &serde_json::Value,
        policy_ids: Option<&Vec<String>>,
        attester_type: &str,
        token_fmt: &str,
        user_id: &str,
    ) -> Result<(Vec<bool>, Vec<PolicyInfo>), AttestationError> {
        let policy_id_list = self.build_policy_id_list(policy_ids, attester_type, user_id).await?;
        
        self.validate_policy_constraints(token_fmt, policy_id_list.len())?;
        
        let (verify_results, evaluate_results) =
            self.evaluate_custom_policies(verify_evidence, &policy_id_list).await?;
        Ok((verify_results, evaluate_results))
    }

    /// Evaluate custom policies
    async fn evaluate_custom_policies(
        &self,
        verify_evidence: &serde_json::Value,
        policy_ids: &[String],
    ) -> Result<(Vec<bool>, Vec<PolicyInfo>), AttestationError> {
        info!("Evaluating custom policies, count: {}", policy_ids.len());
        let mut verify_results = Vec::new();
        let mut evaluate_results = Vec::new();

        let db_connection = get_connection().await.unwrap();
        let policies = match get_policy_by_ids(&db_connection, policy_ids.to_vec()).await {
            Ok(policies) => policies,
            Err(e) => {
                error!("Failed to get policies by ids: {}", e);
                return Err(AttestationError::PolicyNotFoundError(e.to_string()));
            },
        };

        for policy in policies {
            let evaluate_result = match evaluate_policy(verify_evidence, &policy.content) {
                Ok(result) => {
                    let policy_matched = result.get("policy_matched").and_then(|v| v.as_bool()).ok_or_else(|| {
                        let err_msg = format!(
                            "Failed to extract boolean 'policy_matched' from policy evaluation result for policy {}",
                            policy.id
                        );
                        error!("{}", err_msg);
                        AttestationError::PolicyVerificationError(err_msg)
                    })?;

                    verify_results.push(policy_matched);
                    PolicyInfo {
                        appraisal_policy_id: policy.id.clone(),
                        policy_version: policy.version,
                        policy_matched,
                        custom_data: result.get("custom_data").map(|v| v.clone()),
                        policy_content: policy.content.clone(),
                    }
                },
                Err(e) => {
                    error!("Policy evaluation failed for policy {}: {}", policy.id, e);
                    return Err(AttestationError::PolicyVerificationError(format!(
                        "Policy evaluation failed for policy {}: {}",
                        policy.id, e
                    )));
                },
            };
            evaluate_results.push(evaluate_result);
        }

        Ok((verify_results, evaluate_results))
    }
}

impl AttestationHandler for CustomPolicyHandler {
    fn handle<'a>(&'a self, context: &'a mut AttestationContext) -> Pin<Box<dyn Future<Output = Result<(), AttestationError>> + Send + 'a>> {
        Box::pin(async move {
            if let Some(verify_evidence) = &context.verify_evidence {
                let (verify_results, evaluate_results) = self.evaluate_user_policies(
                    verify_evidence,
                    context.evidence.policy_ids.as_ref(),
                    &context.evidence.attester_type,
                    context.measurement.token_fmt.as_deref().unwrap_or("eat"),
                    &context.user_id,
                ).await?;
                
                context.custom_verify_results = verify_results;
                context.custom_evaluate_results = evaluate_results;
            }
            Ok(())
        })
    }
    
    fn set_next(&mut self, next: Box<dyn AttestationHandler>) {
        self.next = Some(next);
    }
    
    fn get_next(&self) -> Option<&dyn AttestationHandler> {
        self.next.as_ref().map(|h| h.as_ref())
    }
    
    fn name(&self) -> &'static str {
        "CustomPolicy"
    }
}